#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import rospy
import moveit_commander
import moveit_msgs.msg
import geometry_msgs.msg
import numpy as np
from sensor_msgs.msg import JointState
import time

class TrajectoryRecorder(object):
    def __init__(self):
        self.arm_group = moveit_commander.MoveGroupCommander("arm")
        self.trajectory = []
        # 订阅关节状态，调整队列大小和频率
        self.joint_states_sub = rospy.Subscriber('/my_gen3/joint_states', JointState, self.joint_states_callback, queue_size=1)
        self.sample_rate = 1  # 采样频率为1Hz
        self.sample_interval = 1.0 / self.sample_rate
        self.last_sample_time = rospy.get_time()

    def joint_states_callback(self, msg):
        current_time = rospy.get_time()
        if current_time - self.last_sample_time >= self.sample_interval:
            # 记录关节位置
            self.trajectory.append(msg.position)
            self.last_sample_time = current_time

    def save_trajectory(self, directory):
        if not os.path.exists(directory):
            os.makedirs(directory)
        date_today = time.strftime('%y%m%d')
        filename = os.path.join(directory, date_today)
        count = 1
        while os.path.exists(filename + '_' + str(count) + '.npy'):
            count += 1
        np.save(filename + '_' + str(count) + '.npy', np.array(self.trajectory))

class TrajectoryPlayer(object):
    def __init__(self):
        self.arm_group = moveit_commander.MoveGroupCommander("arm")

    def execute_trajectory(self, filename):
        try:
            trajectory = np.load(filename)
        except IOError:
            rospy.logerr("加载轨迹文件失败。")
            return False

        for positions in trajectory:
            # 设置目标关节位置
            self.arm_group.set_joint_value_target(positions)
            # 规划轨迹
            plan = self.arm_group.plan()
            # 执行轨迹
            if not self.arm_group.execute(plan, wait=True):
                rospy.logerr("执行轨迹失败。")
                return False

        rospy.loginfo("轨迹播放完成。")
        return True

def record_trajectory():
    rospy.loginfo("开始记录轨迹，按数字0结束并保存。")
    recorder = TrajectoryRecorder()
    while True:
        choice = raw_input("输入数字0结束轨迹记录并保存：")
        if choice == '0':
            rospy.loginfo("结束轨迹记录，保存轨迹...")
            recorder.save_trajectory('/home/liuliyang/catkin_workspace/src/cv_gen3_driver/recordjointstate')
            break

def play_trajectory():
    rospy.loginfo("播放保存的轨迹...")
    directory = '/home/liuliyang/catkin_workspace/src/cv_gen3_driver/recordjointstate'
    files = sorted(os.listdir(directory))
    files = [f for f in files if f.endswith('.npy')]
    if not files:
        rospy.loginfo("没有找到保存的轨迹文件。")
        return

    for i, filename in enumerate(files, start=1):
        print("{}: {}".format(i, filename))

    choice = raw_input("选择要播放的轨迹文件编号：")
    if not choice.isdigit() or int(choice) < 1 or int(choice) > len(files):
        rospy.logwarn("无效的选择，请重新输入。")
        return

    filename = os.path.join(directory, files[int(choice) - 1])
    player = TrajectoryPlayer()
    player.execute_trajectory(filename)

def main():
    rospy.init_node('arm_trace_record_replay')
    moveit_commander.roscpp_initialize(sys.argv)
    while not rospy.is_shutdown():
        choice = raw_input("输入数字进行操作：1-记录轨迹，2-播放保存的轨迹，0-退出：")
        if choice == '1':
            record_trajectory()
        elif choice == '2':
            play_trajectory()
        elif choice == '0':
            rospy.loginfo("退出程序。")
            break
        else:
            rospy.logwarn("无效的选择，请重新输入。")
    moveit_commander.roscpp_shutdown()

if __name__ == '__main__':
    main()
